﻿namespace SpikingAnalyticsLib

open System
open System.IO
open System.Linq    // for viewing List<FiringEvent> as IEnumerable
open System.Collections.Generic
open SpikingNeuronLib

//=================================================================================================

// Extension methods on the SpikingNeuronLib.Triplet type
module TripletExtensions =

    type SpikingNeuronLib.Triplet with

    //---------------------------------------------------------------------------------------------------------
    // Instance Extension methods

    /// Find sets of matching triplets for each event in the specified triplet
    // Use with Triplet.GenerateNeuronTripletMap()
    // See also FindFuzzyMatches() and FindFuzzyMatchesAndRejects()
    member this.FindFuzzyMatchedEvents(tripletMap:Map<int, (Set<int> * Triplet) list>) =
        this.Events
        |> Seq.map (fun ev -> ev.AsTuple)
        |> Seq.map (fun (time, nindex) ->
            // for each event, generate a set of fuzzy-matched triplets
            if tripletMap.ContainsKey nindex then
                tripletMap.[nindex]
                |> Seq.filter (fun (fuzzyTimes, sourceTriplet) -> fuzzyTimes.Contains time)
                |> Seq.map (fun (fuzzyTimes, sourceTriplet) -> sourceTriplet)
                |> Set.ofSeq
            else
                Set.empty
        )
        |> Seq.toArray

    /// Find triplets that fuzzy match this triplet on two or more events
    member this.FindFuzzyMatches(tripletMap:Map<int, (Set<int> * Triplet) list>) =

        let matchedEvents = this.FindFuzzyMatchedEvents(tripletMap)

        // if any two of these sets overlap then the triplet is matched
        seq {
            yield Set.intersect matchedEvents.[0] matchedEvents.[1]
            yield Set.intersect matchedEvents.[0] matchedEvents.[2];
            yield Set.intersect matchedEvents.[1] matchedEvents.[2];
        }
        |> Seq.filter (fun intersection -> intersection.Count > 0)
        |> Seq.concat
        |> Set.ofSeq

    /// Find triplets that fuzzy match this triplet on two or more events
    /// Return both matches and rejects (i.e. single event matches)
    member this.FindFuzzyMatchesAndRejects(tripletMap:Map<int, (Set<int> * Triplet) list>) =

        let matchedEvents = this.FindFuzzyMatchedEvents(tripletMap)

        // if any two of these sets overlap then the triplet is matched
        let matched =
            seq {
                yield Set.intersect matchedEvents.[0] matchedEvents.[1]
                yield Set.intersect matchedEvents.[0] matchedEvents.[2];
                yield Set.intersect matchedEvents.[1] matchedEvents.[2];
            }
            |> Seq.filter (fun intersection -> intersection.Count > 0)
            |> Seq.concat
            |> Set.ofSeq

        let rejected =
            matchedEvents
            |> Seq.concat
            |> Seq.filter (fun triplet -> not (matched.Contains triplet))
            |> Set.ofSeq

        matched, rejected

    /// Report on matches and rejects with the specified triplet map
    member this.CreateMatchReport(tripletMap:Map<int, (Set<int> * Triplet) list>) =
        printfn "Source: %A" this
        let matches, rejected = this.FindFuzzyMatchesAndRejects(tripletMap)
        matches
        |> Seq.iter (fun triplet -> printfn " Match: %A" triplet)
        rejected
        |> Seq.iter (fun triplet -> printfn "Reject: %A" triplet)

    /// Report on matches and rejects with the specified group of triplets
    member this.CreateMatchReport(sourceTriplets:#seq<Triplet>, jitter) =

        let tripletMap = Triplet.GenerateNeuronTripletMap(sourceTriplets, jitter)
        this.CreateMatchReport(tripletMap)

    /// Generate a set of firing times with jitter for all neurons in the triplet
    member this.GenerateNeuronFiringTimeMap(jitter) =
        this.Events.AsEnumerable()
        |> Seq.map (fun ev -> ev.AsTuple)
        |> Seq.fold (fun (state:Map<int, int list>) (time, nindex) ->
            let firingTimes =
                if state.ContainsKey nindex then
                    state.[nindex]
                else
                    []
            Map.add nindex (time :: firingTimes) state
        ) Map.empty
        |> Map.map (fun nindex firingTimes ->
            firingTimes
            |> Seq.map (fun offset ->
                set [ for i in offset - jitter..offset + jitter -> i ]
            )
            |> Set.unionMany        
        )

    // Jitter the firing time of each firing event by 1 to jitter msecs (inclusive)
    member this.Jitter(jitter) =
        let rnd = Random(int(DateTime.Now.ToFileTimeUtc()))

        // the generated triplet must not be equivalent (same transpose) to the original
        // keep generating triplets until this condition is met
        Seq.initInfinite (fun index ->
            let times =
                this.FiringTimes
                |> Seq.map (fun time -> time, rnd.Next(0, jitter) + 1, rnd.NextDouble())
                |> Seq.map (fun (time, magnitude, directionIndicator) ->
                    if directionIndicator > 0.5 then
                        time + magnitude
                    else
                        time - magnitude
                )

            let events =
                Seq.zip times this.Neurons
                |> Seq.map (fun (time, nindex) -> new FiringEvent(time, nindex, EventLabel.Foreground))
                |> Seq.toArray
            (new Triplet(events)).Transpose()
        )
        |> Seq.filter (fun jitteredTriplet -> not (jitteredTriplet.IsEquivalent(this)))
        |> Seq.take 1
        |> Seq.exactlyOne

    //---------------------------------------------------------------------------------------------------------
    // Static Extension methods

    /// Generate the specified number of random triplets where:
    ///  1. each firing event has a unique neuron
    ///  2. none of the generated triplets are isomorphic with any of the triplets in the specified exclude set
    ///  3.triplet events which fuzzy match two or more events from a PNG triplet are rejected
    // see FindFuzzyMatches()
    // Triplets are isomorphic if they have the same shape when their firing events are made co-temporaneous
    // i.e. when both triplets are transposed (temporally shifted to baseline), they have the same firing events
    // (see Triplet.IsEquivalent)
    static member GenerateRandomTriplets(numTriplets, excludeSet:Set<Triplet>) =

        let GenerateTriplet (rnd:Random) =
            let GenerateSampleWithReplacement inclusiveMinimum exclusiveMaximum =
                Seq.initInfinite (fun index ->
                    rnd.Next(inclusiveMinimum, exclusiveMaximum)
                )

            Seq.initInfinite (fun index ->
                let tripletEvents =
                    Seq.zip (GenerateSampleWithReplacement 0 20) (GenerateSampleWithReplacement 0 800)     // excitatory only
                    |> Seq.take 3
                    |> Seq.map (fun (time, nindex) -> new FiringEvent(time, nindex, EventLabel.Foreground))
                    |> Seq.toArray
                new Triplet(tripletEvents)
            )
            // the triplet must not be isomorphic with a triplet in the exclude set
            |> Seq.map (fun triplet -> triplet.Transpose())
            |> Seq.filter (fun triplet -> not (excludeSet.Contains triplet))
            // all three neurons must be different
            |> Seq.filter (fun triplet ->
                let neuronSet = Set.ofSeq triplet.Neurons
                neuronSet.Count = 3
            )
            |> Seq.take 1
            |> Seq.exactlyOne

        let rnd = new Random(int(DateTime.Now.ToFileTimeUtc()))

        // keep generating triplets until the set of unique triplets meets the required number
        Seq.unfold (fun (state:Set<Triplet>) ->
            if state.Count = numTriplets then
                None
            else
                let triplet = GenerateTriplet rnd
                Some(triplet, state.Add triplet)
        ) Set.empty
        |> Seq.toArray

    /// Generate a list of fuzzy firing times and matching Triplets keyed by neuron
    // Each item in the list is a pair of fuzzy time and triplet
    // the fuzzy time is derived from the firing time of the neuron key within that triplet
    static member GenerateNeuronTripletMap(triplets:#seq<Triplet>, jitter) =
        triplets
        |> Seq.map (fun triplet ->
            triplet.Events.AsEnumerable()
            |> Seq.map (fun ev ->
                let fuzzyTime =
                    ev.GetFuzzyFiringTime(jitter)
                    |> Set.ofSeq
                ev.NeuronIndex, (fuzzyTime, triplet)
            )
        )
        |> Seq.concat
        |> Seq.fold (fun (state:Map<int, (Set<int> * Triplet) list>) (nindex, data) ->
            let matchList =
                if state.ContainsKey nindex then
                    state.[nindex]
                else
                    []
            Map.add nindex (data :: matchList) state
        ) Map.empty

    /// Report on matches and rejects for the specified triplets with another group of triplets
    static member CreateMatchReports(triplets:#seq<Triplet>, sourceTriplets:#seq<Triplet>, jitter) =
        let tripletMap = Triplet.GenerateNeuronTripletMap(sourceTriplets, jitter)

        triplets
        |> Seq.iter (fun triplet ->
            triplet.CreateMatchReport(tripletMap)
            printfn ""
        )

    /// Filter out triplets that fuzzy match the specified triplets with the specified source group
    // Curried arguments for convenience
    static member FilterRejectFuzzyMatchedTriplets (sourceTriplets:#seq<Triplet>) jitter (triplets:#seq<Triplet>) =
        let tripletMap = Triplet.GenerateNeuronTripletMap(sourceTriplets, jitter)

        triplets
        |> Seq.filter (fun triplet ->
            let matches = triplet.FindFuzzyMatches(tripletMap)
            matches.Count = 0
        )
        |> Seq.toArray

    /// Filter the triplets, keeping only those that fuzzy match the specified triplets with the specified source group
    // Curried arguments for convenience
    static member FilterKeepOnlyFuzzyMatchedTriplets (sourceTriplets:#seq<Triplet>) jitter (triplets:#seq<Triplet>) =
        let tripletMap = Triplet.GenerateNeuronTripletMap(sourceTriplets, jitter)

        triplets
        |> Seq.filter (fun triplet ->
            let matches = triplet.FindFuzzyMatches(tripletMap)
            matches.Count > 0
        )
        |> Seq.toArray
